import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import spearmanr
from pathlib import Path
import joblib
import xgboost as xgb
from featurework import prepare_data, detect_encoding
from sklearn.metrics import precision_score, recall_score, f1_score
from matplotlib_venn import venn2  # 需安装matplotlib-venn: pip install matplotlib-venn

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 全局配置
PROJECT_ROOT = Path(__file__).resolve().parents[2]
MODEL_DIR = PROJECT_ROOT / "model"
DATA_DIR = PROJECT_ROOT / "data"
PLOT_DIR = PROJECT_ROOT / "plots"
PLOT_DIR.mkdir(parents=True, exist_ok=True)  # 创建图形保存目录


# -------------------------- 1. 数据加载与评估函数 --------------------------
def load_check_data(check_path: Path) -> tuple[list, list]:
    """加载实际结果"""
    check_encoding = detect_encoding(check_path)
    check_df = pd.read_csv(check_path, encoding=check_encoding)
    actual_top10 = check_df["涨幅最大股票代码"].dropna().tolist()
    actual_bottom10 = check_df["涨幅最小股票代码"].dropna().tolist()
    return actual_top10, actual_bottom10


def calculate_metrics(
        pred_top10: list, pred_bottom10: list,
        actual_top10: list, actual_bottom10: list
) -> dict:
    """计算评估指标"""
    # 计算F1分数（涨幅最大）
    pred_top_label = [1 if code in pred_top10 else 0 for code in actual_top10]
    actual_top_label = [1] * len(actual_top10)
    precision_up = precision_score(actual_top_label, pred_top_label, zero_division=0)
    recall_up = recall_score(actual_top_label, pred_top_label, zero_division=0)
    f1_up = f1_score(actual_top_label, pred_top_label, zero_division=0)

    # 计算F1分数（涨幅最小）
    pred_bottom_label = [1 if code in pred_bottom10 else 0 for code in actual_bottom10]
    actual_bottom_label = [1] * len(actual_bottom10)
    precision_down = precision_score(actual_bottom_label, pred_bottom_label, zero_division=0)
    recall_down = recall_score(actual_bottom_label, pred_bottom_label, zero_division=0)
    f1_down = f1_score(actual_bottom_label, pred_bottom_label, zero_division=0)

    # 计算Spearman系数
    pred_rank = {code: idx + 1 for idx, code in enumerate(pred_top10)}
    pred_rank.update({code: idx + 11 for idx, code in enumerate(pred_bottom10)})
    actual_rank = {code: idx + 1 for idx, code in enumerate(actual_top10)}
    actual_rank.update({code: idx + 11 for idx, code in enumerate(actual_bottom10)})

    common_codes = list(set(pred_rank.keys()) & set(actual_rank.keys()))[:10]
    pred_ranks = [pred_rank[code] for code in common_codes]
    actual_ranks = [actual_rank[code] for code in common_codes]

    spearman_up, _ = spearmanr(pred_ranks[:len(actual_top10)], actual_ranks[:len(actual_top10)])
    spearman_down, _ = spearmanr(pred_ranks[len(actual_top10):], actual_ranks[len(actual_top10):])
    spearman_up = round(spearman_up if not np.isnan(spearman_up) else 0, 3)
    spearman_down = round(spearman_down if not np.isnan(spearman_down) else 0, 3)

    # 最终得分
    final_score = round(0.2 * f1_up + 0.2 * f1_down + 0.3 * spearman_up + 0.3 * spearman_down, 3)

    return {
        "precision_up": precision_up, "recall_up": recall_up, "f1_up": f1_up,
        "precision_down": precision_down, "recall_down": recall_down, "f1_down": f1_down,
        "spearman_up": spearman_up, "spearman_down": spearman_down,
        "final_score": final_score,
        "pred_top10": pred_top10, "pred_bottom10": pred_bottom10
    }


# -------------------------- 2. 模型加载与预测 --------------------------
def load_models():
    """加载模型"""
    lgb_models = [joblib.load(f) for f in MODEL_DIR.glob("lgb_model_fold*.pkl")]
    xgb_models = [joblib.load(f) for f in MODEL_DIR.glob("xgb_model_fold*.pkl")]
    return lgb_models, xgb_models


def ensemble_predict(lgb_models, xgb_models, X_test, lgb_weight):
    """集成预测"""
    lgb_pred = np.mean([m.predict(X_test) for m in lgb_models], axis=0)
    xgb_pred = np.mean([m.predict(xgb.DMatrix(X_test)) for m in xgb_models], axis=0)
    return lgb_weight * lgb_pred + (1 - lgb_weight) * xgb_pred


# -------------------------- 3. 图形绘制函数 --------------------------
def plot_weight_analysis(metrics_df, best_weight, actual_top10, actual_bottom10):
    """绘制权重分析图形"""
    # 创建2x2图形布局
    fig, axes = plt.subplots(2, 2, figsize=(16, 14))
    fig.suptitle("模型权重与性能分析", fontsize=16, y=0.95)

    # 1. 权重与最终得分关系图
    ax1 = axes[0, 0]
    ax1.plot(metrics_df["lgb_weight"], metrics_df["final_score"],
             marker='o', color='#2ecc71', linewidth=2, markersize=6)
    ax1.axvline(x=best_weight, color='r', linestyle='--',
                label=f'最优权重: {best_weight:.2f}')
    ax1.set_title("权重组合与最终得分关系", fontsize=12)
    ax1.set_xlabel("LightGBM权重 (XGBoost权重=1-此值)")
    ax1.set_ylabel("最终得分")
    ax1.grid(alpha=0.3)
    ax1.legend()

    # 2. F1分数对比图
    ax2 = axes[0, 1]
    x = np.arange(len(metrics_df))
    width = 0.35
    ax2.bar(x - width / 2, metrics_df["f1_up"], width, label='涨幅最大股票', color='#3498db')
    ax2.bar(x + width / 2, metrics_df["f1_down"], width, label='涨幅最小股票', color='#9b59b6')
    ax2.axvline(x=metrics_df[metrics_df["lgb_weight"] == best_weight].index[0],
                color='r', linestyle='--')
    ax2.set_title("不同权重下的F1分数对比", fontsize=12)
    ax2.set_xlabel("权重序号")
    ax2.set_ylabel("F1分数")
    ax2.set_xticks(x[::2])  # 每隔一个刻度显示
    ax2.set_xticklabels([f"{w:.1f}" for w in metrics_df["lgb_weight"][::2]])
    ax2.legend()

    # 3. Spearman系数对比图
    ax3 = axes[1, 0]
    ax3.bar(x - width / 2, metrics_df["spearman_up"], width, label='涨幅最大股票', color='#e74c3c')
    ax3.bar(x + width / 2, metrics_df["spearman_down"], width, label='涨幅最小股票', color='#f39c12')
    ax3.axvline(x=metrics_df[metrics_df["lgb_weight"] == best_weight].index[0],
                color='r', linestyle='--')
    ax3.set_title("不同权重下的Spearman系数对比", fontsize=12)
    ax3.set_xlabel("权重序号")
    ax3.set_ylabel("Spearman系数")
    ax3.set_xticks(x[::2])
    ax3.set_xticklabels([f"{w:.1f}" for w in metrics_df["lgb_weight"][::2]])
    ax3.legend()

    # 4. 最优权重下的预测与实际重叠韦恩图
    ax4 = axes[1, 1]
    best_idx = metrics_df[metrics_df["lgb_weight"] == best_weight].index[0]
    best_pred_top10 = metrics_df.iloc[best_idx]["pred_top10"]

    # 绘制涨幅最大股票的韦恩图
    venn2([set(best_pred_top10), set(actual_top10)],
          set_labels=('预测涨幅最大', '实际涨幅最大'), ax=ax4)
    ax4.set_title(f"最优权重下预测与实际重叠情况\n(LightGBM权重={best_weight:.2f})", fontsize=12)

    # 调整布局并保存
    plt.tight_layout()
    plt.savefig(PLOT_DIR / "weight_analysis.png", dpi=300, bbox_inches="tight")
    plt.close()
    print(f"\n图形已保存至: {PLOT_DIR / 'weight_analysis.png'}")


# -------------------------- 4. 主函数：最优权重搜索与可视化 --------------------------
def search_best_weight(weight_steps=0.1):
    """搜索最优权重并可视化结果"""
    # 准备数据
    _, _, X_test, stock_codes = prepare_data(predict=True)
    actual_top10, actual_bottom10 = load_check_data(DATA_DIR / "check.csv")
    lgb_models, xgb_models = load_models()

    # 遍历权重
    weights = np.arange(0, 1.01, weight_steps)
    metrics_list = []
    best_score = -1
    best_weight = 0.5

    print("\n===== 开始搜索最优权重 =====")
    for lgb_w in weights:
        xgb_w = 1 - lgb_w
        print(f"\n----- 权重组合：LightGBM={lgb_w:.1f}, XGBoost={xgb_w:.1f} -----")

        # 预测与筛选
        pred = ensemble_predict(lgb_models, xgb_models, X_test, lgb_w)
        pred_df = pd.DataFrame({"股票代码": stock_codes, "预测涨跌幅": pred})
        pred_df_sorted = pred_df.sort_values("预测涨跌幅", ascending=False)
        pred_top10 = pred_df_sorted.head(10)["股票代码"].tolist()
        pred_bottom10 = pred_df_sorted.tail(10)["股票代码"].tolist()

        # 计算指标
        metrics = calculate_metrics(pred_top10, pred_bottom10, actual_top10, actual_bottom10)
        metrics["lgb_weight"] = lgb_w
        metrics["xgb_weight"] = xgb_w
        metrics_list.append(metrics)

        # 打印当前结果
        print(f"预测涨幅最大Top10: {pred_top10[:5]}...")  # 简化显示
        print(f"预测涨幅最小Top10: {pred_bottom10[:5]}...")
        print(f"最终得分: {metrics['final_score']:.3f}")

        # 更新最优结果
        if metrics["final_score"] > best_score:
            best_score = metrics["final_score"]
            best_weight = lgb_w

    # 整理结果为DataFrame
    metrics_df = pd.DataFrame(metrics_list)

    # 绘制图形
    plot_weight_analysis(metrics_df, best_weight, actual_top10, actual_bottom10)

    # 终端显示最优结果
    print("\n" + "=" * 50)
    print("===== 最优权重搜索结果 =====")
    print(f"最优LightGBM权重: {best_weight:.2f}")
    print(f"最优XGBoost权重: {1 - best_weight:.2f}")
    print(f"最优最终得分: {best_score:.3f}")
    print("=" * 50)


if __name__ == "__main__":
    search_best_weight(weight_steps=0.1)  # 步长0.1，共11组权重
